from parse.token import Token, TokenType


RESERVED_ID_DICT = {
    '.': TokenType.DOT,
    'class': TokenType.CLASS,
    'else': TokenType.ELSE,
    'false': TokenType.FALSE,
    'true': TokenType.TRUE,
    'fun': TokenType.FUN,
    'for': TokenType.FOR,
    'if': TokenType.IF,
    'nil': TokenType.NIL,
    'or': TokenType.OR,
    'print': TokenType.PRINT,
    'return': TokenType.RETURN,
    'super': TokenType.SUPER,
    'this': TokenType.THIS,
    'var': TokenType.VAR,
    'while': TokenType.WHILE,
    'break': TokenType.BREAK,
    'continue': TokenType.CONTINUE,
}


class Lexer:
    def __init__(self):
        self.line = 1
        self.source = ""
        self.tokens = []

    def parse(self, s: str):
        self.source = s

        while not self.is_eof():
            c = self.advance()
            match c:
                case '(':
                    self.add_token(c, TokenType.LEFT_PARENT)
                case ')':
                    self.add_token(c, TokenType.RIGHT_PARENT)
                case '{':
                    self.add_token(c, TokenType.LEFT_BRACE)
                case '}':
                    self.add_token(c, TokenType.RIGHT_BRACE)
                case ',':
                    self.add_token(c, TokenType.COMMA)
                case '.':
                    self.add_token(c, TokenType.DOT)
                case '-':
                    self.add_token(c, TokenType.MINUS)
                case '+':
                    self.add_token(c, TokenType.PLUS)
                case ';':
                    self.add_token(c, TokenType.SEMICOLON)
                case '/':
                    if self.peek() == '/':
                        self.advance()
                        self.skip_comment()
                    else:
                        self.add_token(c, TokenType.SLASH)
                case '*':
                    self.add_token(c, TokenType.STAR)
                case '!':
                    if self.peek() == '=':
                        self.add_token('!=', TokenType.BANG_EQUAL)
                        self.advance()
                    else:
                        self.add_token(c, TokenType.BANG)
                case '=':
                    if self.peek() == '=':
                        self.add_token('==', TokenType.EQUAL_EQUAL)
                    else:
                        self.add_token(c, TokenType.EQUAL)
                case '>':
                    if self.peek() == '=':
                        self.add_token('>=', TokenType.GREATER_EQUAL)
                        self.advance()
                    else:
                        self.add_token(c, TokenType.GREATER)
                case '<':
                    if self.peek() == '=':
                        self.add_token('<=', TokenType.LESS_EQUAL)
                    else:
                        self.add_token('<', TokenType.LESS)
                case '"' | "'":
                    self.add_token(self.parse_string(c), TokenType.STRING)
                case n if n.isdigit():
                    self.add_token(self.parse_number(n), TokenType.NUMBER)
                case s if s.isalpha() or s == '_':
                    lexeme = self.parse_identifier(s)
                    token_type = RESERVED_ID_DICT.get(lexeme, TokenType.IDENTIFIER)
                    self.add_token(lexeme, token_type)
                case '\n':
                    self.line += 1
                case sp if sp.isspace():
                    pass
                case _:
                    raise ValueError(f'invalid char "{c}"')

        self.add_token('', TokenType.EOF)

        return self.tokens

    def add_token(self, lexeme: str, token_type: TokenType):
        self.tokens.append(Token(lexeme, token_type, self.line))

    def peek(self) -> str:
        if not self.is_eof():
            return self.source[0]
        else:
            return ''

    def is_eof(self):
        return not self.source

    def peek_next(self):
        if len(self.source) == 1:
            return ''
        return self.source[1]

    def advance(self):
        c = self.source[0]
        self.source = self.source[1:]
        return c

    def parse_string(self, wrap_str: str):
        s_list = []
        
        while self.peek() != wrap_str and (not self.is_eof()):
            s_list.append(self.advance())

        if self.is_eof():
            raise ValueError('unclosed string')

        self.advance()

        return ''.join(s_list)

    def parse_number(self, head: str):
        number_list = [head]

        while self.peek().isdigit() and (not self.is_eof()):
            number_list.append(self.advance())

        if self.peek() == '.':
            while self.peek().isdigit() and (not self.is_eof()):
                number_list.append(self.advance())

        return ''.join(number_list)

    def parse_identifier(self, head: str):
        id_list = [head]
        while self.peek().isalnum() and (not self.is_eof()):
            id_list.append(self.advance())

        return ''.join(id_list)

    def skip_comment(self):
        while not self.is_eof() and not self.peek() == '\n':
            self.advance()

